import time
import numpy as np

import torch
import torch.nn as nn


def model_train(model, trainloader, optimizer, epoch, criterion, args):
    model.train()
    acc = 0
    if (epoch % 10 == 0) and epoch < 400:
        args.lr /= 10
        optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)

    # ts = time.time()
    out_time = time.time()
    for i_iter, batch in enumerate(trainloader):
        images, labels, fileNames = batch
        images = images.cuda()
        labels = labels.cuda()

        optimizer.zero_grad()

        # ts = time.time()
        outs = model(images)
        # print('model: %.1f s' % (time.time() - ts))

        loss = criterion(outs, labels)
        # ts = time.time()

        loss.backward()
        # print('backward: %.1f s' % (time.time() - ts))

        # ts = time.time()
        optimizer.step()
        # print('optimizer.step: %.1f s' % (time.time() - ts))

        _, predicted = torch.max(outs.data, 1)
        # 累加识别正确的样本数
        acc += (predicted == labels).sum()

        # print(conf)
        if args.local_rank == 1 and (i_iter + 1) % args.log_interval == 0:
            print('epoch_batch: {:d}_{:d} | train_loss: {:f}  | : train_acc: {:f} | time:  {:f} s'.format(
                epoch, i_iter, loss.item(), float(acc) / ((i_iter + 1) * args.batch_size), (time.time() - out_time)))
            out_time = time.time()


def model_test(model, testloader, epoch, args):
    model.eval()
    test_acc = 0

    # if args.local_rank == 0:
    test_time = time.time()
    #     conf0 = np.zeros((5, 5))
    # else:
    conf1 = np.zeros((5, 5))
    for i_iter, batch in enumerate(testloader):
        images, labels, fileNames = batch
        images = images.cuda()
        labels = labels.cuda()

        # ts = time.time()
        outs = model(images)
        # print('model: %.1f s' % (time.time() - ts))

        # 累加识别正确的样本数
        _, predicted = torch.max(outs.data, 1)
        test_acc += (predicted == labels).sum()

        predicted_cpu = predicted.cpu()
        labels_cpu = labels.cpu()
        # if args.local_rank == 0:
        #     for i in range(args.batch_size):
        #         conf0[labels_cpu[i]][predicted_cpu[i]] += 1
        # else:
        for j in range(args.batch_size):
            conf1[labels_cpu[j]][predicted_cpu[j]] += 1

    # print(conf)
    # if args.local_rank == 0:
    #     print(conf0)
    # else:
    print(conf1)

    if args.local_rank == 1:
        print('epoch: {:d} | : test_acc: {:f} | test_time:  {:f} s'.format(
            epoch, float(test_acc) / ((i_iter + 1) * args.batch_size), (time.time() - test_time)))


class CenterLoss(nn.Module):
    """Center loss.
    Reference:
    Wen et al. A Discriminative Feature Learning Approach for Deep Face Recognition. ECCV 2016.
    Args:
        num_classes (int): number of classes.
        feat_dim (int): feature dimension.
    """
    def __init__(self, num_classes=5, feat_dim=32, use_gpu=True):
        super(CenterLoss, self).__init__()
        self.num_classes = num_classes
        self.feat_dim = feat_dim
        self.use_gpu = use_gpu
        if self.use_gpu:
            self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim).cuda())
        else:
            self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim))
    def forward(self, x, labels):
        """
        Args:
            x: feature matrix with shape (batch_size, feat_dim).
            labels: ground truth labels with shape (num_classes).
        """
        assert x.size(0) == labels.size(0), "features.size(0) is not equal to labels.size(0)"
        batch_size = x.size(0)
        distmat = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(batch_size, self.num_classes) + \
                  torch.pow(self.centers, 2).sum(dim=1, keepdim=True).expand(self.num_classes, batch_size).t()
        distmat.addmm_(1, -2, x, self.centers.t())
        classes = torch.arange(self.num_classes).long()
        if self.use_gpu: classes = classes.cuda()
        labels = labels.unsqueeze(1).expand(batch_size, self.num_classes)
        mask = labels.eq(classes.expand(batch_size, self.num_classes))
        # print(mask)
        dist = []
        for i in range(batch_size):
            # print(mask[i])
            value = distmat[i][mask[i]]
            value = value.clamp(min=1e-12, max=1e+12)  # for numerical stability
            dist.append(value)
        dist = torch.cat(dist)
        loss = dist.mean()
        return loss

